from typing import Callable, TYPE_CHECKING, Union

from zkl_aiutils_datasets.processing import KnownSizeChainedDatasets
from .parquet_file_row_group import ParquetFileRowGroupDataset
from .source import parquet_file_source_context

if TYPE_CHECKING:
    from pyarrow.parquet import ParquetFile, FileMetaData


class ParquetFileDataset(KnownSizeChainedDatasets[dict]):
    def __init__(self, source: Union[Callable[[], 'ParquetFile'], 'ParquetFile']):
        with parquet_file_source_context(source) as parquet_file:
            metadata = parquet_file.metadata
        super().__init__([
            ParquetFileRowGroupDataset(source, row_group_index, metadata.row_group(row_group_index))
            for row_group_index in range(metadata.num_row_groups)])
        self._metadata = metadata

    @property
    def metadata(self) -> 'FileMetaData':
        return self._metadata
